#---------------------------------------------------------------#
#    ResNet 
#---------------------------------------------------------------# 

import tensorflow as tf
import tensorflow.keras as keras

import numpy as ny

'''
'''


def BasicBlock(x, channel, strides=(1,1), expansion=1):
    resdual = keras.layers.Conv2D(channel, (3,3), strides=strides, padding='same', use_bias=False)(x)
    resdual = keras.layers.BatchNormalization()(resdual)
    resdual = keras.layers.Activation('relu')(resdual)
    resdual = keras.layers.Conv2D(channel*expansion, (3,3), padding='same', use_bias=False)(resdual)

    shortcut = x
    if strides != (1,1) or x.shape[-1] != channel*expansion:
        shortcut = keras.layers.Conv2D(channel*expansion, (1,1), strides=strides, use_bias=False)(shortcut)
        shortcut = keras.layers.BatchNormalization()(shortcut)

    y = resdual + shortcut
    y = keras.layers.Activation('relu')(y)

    return y

def BottleNeck(x, channel, strides=(1,1), expansion=4):
    residual = keras.layers.Conv2D(channel, (1,1), use_bias=False)(x)
    residual = keras.layers.BatchNormalization()(residual)
    residual = keras.layers.Activation('relu')(residual)
    residual = keras.layers.Conv2D(channel, (3,3), strides=strides, padding='same', use_bias=False)(residual)
    residual = keras.layers.BatchNormalization()(residual)
    residual = keras.layers.Activation('relu')(residual)
    residual = keras.layers.Conv2D(channel*expansion, (1,1), use_bias=False)(residual)
    residual = keras.layers.BatchNormalization()(residual)

    shortcut = x
    if strides != (1,1) or x.shape[-1] != channel*expansion:
        shortcut = keras.layers.Conv2D(channel*expansion, (1,1), strides=strides, use_bias=False)(shortcut)
        shortcut = keras.layers.BatchNormalization()(shortcut)

    y = resdual + shortcut
    y = keras.layers.Activation('relu')(y)

    return y

def ResNet_process(x, block, num_block, num_classes=99):
    '''
        block   ResNet的block过程，为函数
    '''
    assert len(num_block) == 4, 'The given "num_block" is not proper: {}'.format(num_block)

    # feature_map 1
    f1 = keras.layers.Conv2D(64, (7,7), strides=(2,2), padding='same', use_bias=False)(x)
    f1 = keras.layers.BatchNormalization()(f1)
    f1 = keras.layers.Activation('relu')(f1)

    # feature_map 2
    f2 = keras.layers.MaxPooling2D(pool_size=(3,3), strides=(2,2), padding='same')(f1)
    tmp = [1] * (num_block[0] - 1)
    tmp.insert(0, 1)
    for i in tmp:
        f2 = block(f2, channel=64, strides=(i,i))

    # feature_map 3
    tmp = [1] * (num_block[1] - 1)
    tmp.insert(0, 2)
    for i in tmp:
        f3 = block(f2, channel=128, strides=(i,i))

    # feature_map 4
    tmp = [1] * (num_block[2] - 1)
    tmp.insert(0, 2)
    for i in tmp:
        f4 = block(f3, channel=256, strides=(i,i))

    # feature_map 5
    tmp = [1] * (num_block[3] - 1)
    tmp.insert(0, 2)
    for i in tmp:
        f5 = block(f4, channel=512, strides=(i,i))

    y = keras.layers.AveragePooling2D(pool_size=(1,1))(f5)
    y = keras.layers.Flatten()(y)
    y = keras.layers.Dense(num_classes+1)(y)

    return y # f1, f2, f3, f4, f5, y

def get_net_model(net_input_shape, num_classes, block=BasicBlock, num_block=[2,2,2,2]):
    assert net_input_shape == [224,224,3], 'ResNet输入photo的尺寸需要为[224,224,3]，但提供的为{}'.format(net_input_shape)
    # 1.输入层
    i_put = keras.layers.Input(net_input_shape)

    # 2.backbone
    feature_map = ResNet_process(i_put, block=block, num_block=num_block, num_classes=num_classes)

    # 3.输出层
    o_put = keras.layers.Dense(num_classes+1, activation='softmax')(feature_map)

    # 4.模型实例
    net_model = keras.models.Model(i_put, o_put)

    return net_model

if __name__ == '__main__':
    # 测试搭建的网络
    net_input_shape = [224, 224, 3]
    num_classes     = 1
    resnet_18 = get_net_model(net_input_shape, num_classes, block=BasicBlock, num_block=[2,2,2,2])
    resnet_18.summary()

    # net_input_shape.insert(0, 2)
    # # x = np.random.random(net_input_shape).astype(np.float32) # 增加batch轴
    # x = np.ones(net_input_shape).astype(np.float32) / 2
    # y = net_model(x)
    # y = y.numpy()
    # print('测试网络的前向计算功能：')
    # print('    x.shape = {}'.format(x.shape))
    # print('    y.shape = {}'.format(y.shape))

# def get_net_model(net_input_shape, num_classes):
#     assert net_input_shape == [224,224,3]
#     # 1.输入层
#     i_put = keras.layers.Input(net_input_shape)

#     # 2.backbone
#     feature_map = backbone_process(i_put)

#     # 3.输出层
#     o_put = keras.layers.Dense(num_classes+1, activation='softmax')(feature_map)

#     # 4.模型实例
#     net_model = keras.models.Model(i_put, o_put)

#     return net_model